from langchain_community.chat_models import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory

model = ChatZhipuAI(
    model="glm-4-plus",
    temprature=0.5
)

prompts = ChatPromptTemplate.from_messages([
    MessagesPlaceholder(variable_name="history"),
    ("human", '{input}')
])

store = {}


def get_session_history(session_id: str):
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


chat_with_history = RunnableWithMessageHistory(
    prompts | model,
    get_session_history,
    input_messages_key="input",
    history_messages_key="history"
)

# 定义一个session_id的配置
config = {'configurable': {'session_id': 'abc123'}}

# 第一轮会话
response = chat_with_history.invoke(
    input={"input": "你好，我是兰博，请记住我，下一次称呼我为‘兰博’,每次回答前都需要称呼我"}, config=config)
print(response.content)

# 第二轮会话
response = chat_with_history.invoke(input={"input": "你好，我是谁"}, config=config)
print(response.content)


